const express = require('express')
const server = express()
const compression = require('compression')
require('express-ws')(server, null, {
  wsOptions: {
    maxPayload: 16 * 1024 * 1024 // 16MB最大负载
  }
})

// 中间件配置
server.use(compression())
server.use(express.static('./dist/'))

const { NodeSSH } = require('node-ssh')
const ssh = new NodeSSH()

async function handleConnection(socket) {
  console.log('WebSocket连接已建立')
  let isAlive = true
  let timer = null
  let sshCredentials = null
  let sshCredentialsReceived = false
  let sshStream = null

  socket.setMaxListeners(15)

  // 心跳检测机制
  const checkAlive = () => {
    if (!isAlive) {
      console.log('连接超时，即将关闭')
      cleanupResources()
      return
    }
    isAlive = false
    socket.ping(null, { mask: false, binary: true })
    timer = setTimeout(checkAlive, 600000) // 10分钟超时
  }

  // 资源清理函数
  const cleanupResources = () => {
    if (timer) clearTimeout(timer)
    if (sshStream) sshStream.end()
    if (ssh.isConnected()) ssh.dispose()
    if (socket.readyState === socket.OPEN) socket.close()
  }

  // 消息处理主逻辑
  socket.on('message', async (data) => {
    try {
      if (!sshCredentialsReceived) {
        // 处理SSH凭据
        sshCredentials = JSON.parse(data)
        sshCredentialsReceived = true

        try {
          await ssh.connect({
            host: sshCredentials.ip,
            port: sshCredentials.port,
            username: sshCredentials.username,
            password: sshCredentials.password,
            readyTimeout: 5000
          })

          console.log('SSH连接成功')
          sshStream = await ssh.requestShell()
          console.log('SSH Shell创建成功')

          // SSH数据流处理
          sshStream.on('data', (data) => {
            console.log('接收到SSH数据:', data.toString())
            isAlive = true
            if (socket.readyState === socket.OPEN) {
              socket.send(data.toString(), (err) => {
                if (err) console.error('数据发送失败:', err)
              })
            }
          })
          sshStream.on('close', () => {
            console.log('SSH会话结束')
            cleanupResources()
          })

          // 开始心跳检测
          timer = setTimeout(checkAlive, 600000)
        } catch (sshError) {
          console.error('SSH连接失败:', sshError)
          socket.send(
            JSON.stringify({
              error: 'SSH连接失败',
              details: sshError.message
            })
          )
          cleanupResources()
        }
      } else {
        // 处理终端输入和窗口调整
        isAlive = true
        try {
          const message = JSON.parse(data)
          console.log(message, 77)

          if (message.action === 'resize' && sshStream) {
            sshStream.setWindow(message.rows, message.cols)
          } else if (sshStream) {
            sshStream.write(data)
          }
        } catch (e) {
          if (sshStream) sshStream.write(data)
        }
      }
    } catch (parseError) {
      console.error('消息处理错误:', parseError)
      socket.send(JSON.stringify({ error: '无效的消息格式' }))
      cleanupResources()
    }
  })

  // 错误和关闭处理
  socket.on('error', (err) => {
    console.error('WebSocket错误:', err)
    cleanupResources()
  })

  socket.on('close', () => {
    console.log('WebSocket连接关闭')
    cleanupResources()
  })
}

// WebSocket路由
server.ws('/ssh', (ws, req) => {
  handleConnection(ws)
})

// 启动服务器
server.listen(8080, () => {
  console.log('WebSSH服务器已启动在端口8080')
})
